import os

from dotenv import load_dotenv
from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated
import operator
from langchain_core.messages import HumanMessage, AIMessage
from langchain_openai import ChatOpenAI, AzureChatOpenAI


# 1. 定义状态结构 - 团队成员共享的工作区
class TeamState(TypedDict):
    # 项目需求
    requirements: str
    # 各角色的输出
    product_manager_output: str
    developer_output: str
    tester_output: str
    # 最终结果
    final_result: str
    # 消息历史
    messages: Annotated[list, operator.add]


# 2. 初始化不同专业的LLM（实际中可以用同一个）
load_dotenv()
llm = AzureChatOpenAI(
    api_key=os.getenv("AZURE_OPENAI_API_KEY"),
    azure_endpoint=os.getenv("AZURE_OPENAI_ENDPOINT"),
    azure_deployment=os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME"),
    api_version=os.getenv("AZURE_OPENAI_API_VERSION"),
    temperature=0.7
)

# 3. 定义各个角色的Agent节点
def product_manager_node(state: TeamState):
    """产品经理Agent：分析需求，输出产品方案"""
    print("\n👔 [产品经理] 开始工作...")

    requirements = state["requirements"]
    prompt = f"""
    你是一个资深产品经理。请根据以下需求创建产品方案：

    需求: {requirements}

    请输出包含：
    1. 功能列表
    2. 用户流程图
    3. 主要界面描述

    用清晰的结构化格式回复。
    """

    response = llm.invoke([HumanMessage(content=prompt)])
    return {
        "product_manager_output": response.content,
        "messages": [HumanMessage(content=f"产品经理完成了需求分析")]
    }


def developer_node(state: TeamState):
    """开发工程师Agent：根据产品方案写代码"""
    print("\n💻 [开发工程师] 开始工作...")

    pm_output = state["product_manager_output"]
    prompt = f"""
    你是一个全栈开发工程师。根据产品经理的方案编写核心代码：

    产品方案: {pm_output}

    请输出：
    1. 技术选型建议
    2. 核心API设计
    3. 关键代码片段

    用专业的技术文档格式回复。
    """

    response = llm.invoke([HumanMessage(content=prompt)])
    return {
        "developer_output": response.content,
        "messages": [HumanMessage(content=f"开发工程师完成了编码")]
    }


def tester_node(state: TeamState):
    """测试工程师Agent：根据需求和代码写测试用例"""
    print("\n🐛 [测试工程师] 开始工作...")

    requirements = state["requirements"]
    dev_output = state["developer_output"]
    prompt = f"""
    你是一个专业的测试工程师。

    原始需求: {requirements}
    开发输出: {dev_output}

    请创建：
    1. 测试策略
    2. 关键测试用例
    3. 自动化测试建议

    用专业的测试文档格式回复。
    """

    response = llm.invoke([HumanMessage(content=prompt)])
    return {
        "tester_output": response.content,
        "messages": [HumanMessage(content=f"测试工程师完成了测试计划")]
    }


def manager_node(state: TeamState):
    """项目经理Agent：整合所有输出，生成最终报告"""
    print("\n📊 [项目经理] 整合最终报告...")

    pm_output = state["product_manager_output"]
    dev_output = state["developer_output"]
    tester_output = state["tester_output"]

    prompt = f"""
    你是一个项目经理，请整合团队的工作成果：

    产品方案: {pm_output}

    技术实现: {dev_output}

    测试计划: {tester_output}

    请生成一个完整的项目报告，包括：
    1. 项目概述
    2. 核心成果总结
    3. 下一步行动计划

    报告要专业、结构化。
    """

    response = llm.invoke([HumanMessage(content=prompt)])
    return {
        "final_result": response.content,
        "messages": [HumanMessage(content=f"项目经理完成了最终报告")]
    }


# 4. 构建工作流图
def create_team_workflow():
    # 创建图
    workflow = StateGraph(TeamState)

    # 添加节点
    workflow.add_node("product_manager", product_manager_node)
    workflow.add_node("developer", developer_node)
    workflow.add_node("tester", tester_node)
    workflow.add_node("project_manager", manager_node)

    # 定义工作流：产品经理 -> 开发 -> 测试 -> 项目经理
    workflow.set_entry_point("product_manager")
    workflow.add_edge("product_manager", "developer")
    workflow.add_edge("developer", "tester")
    workflow.add_edge("tester", "project_manager")
    workflow.add_edge("project_manager", END)

    return workflow.compile()


# 5. 运行多Agent团队
def run_team_simulation():
    # 创建团队工作流
    team_app = create_team_workflow()

    # 生成流程图  display(Image(chain.get_graph().draw_mermaid_png()))
    # image = team_app.get_graph().draw_mermaid()
    # print(image)

    # 初始化项目需求
    initial_state = {
        "requirements": "开发一个个人任务管理应用，支持创建任务、设置截止日期、分类和搜索功能",
        "product_manager_output": "",
        "developer_output": "",
        "tester_output": "",
        "final_result": "",
        "messages": []
    }

    print("🚀 启动多Agent团队协作...")
    print(f"项目需求: {initial_state['requirements']}")
    print("-" * 50)

    # 执行工作流
    result = team_app.invoke(initial_state)

    print("\n" + "=" * 50)
    print("🎉 项目完成！最终报告:")
    print("=" * 50)
    print(result["final_result"])

    return result


if __name__ == "__main__":
    run_team_simulation()
